﻿#include "StringManager.h"
#include "StringExactRules.h" //精确匹配字典文件
#include "StringFuzzyRules.h" //模糊匹配字典文件
#include <cstring>
#include <algorithm>
#include <queue>
#include <stdexcept>

// 静态成员初始化
std::unordered_map<std::string, std::string> StringManager::default_utf8_rules_;
std::unordered_map<std::string, std::string> StringManager::append_utf8_rules_;
std::unordered_map<std::wstring, std::wstring> StringManager::default_wide_rules_;
std::unordered_map<std::wstring, std::wstring> StringManager::append_wide_rules_;

// 模糊字典分离
std::vector<StringManager::FuzzyUtf8Entry> StringManager::fuzzy_utf8_default_dict_;
std::vector<StringManager::FuzzyUtf8Entry> StringManager::fuzzy_utf8_append_dict_;
std::vector<StringManager::FuzzyWideEntry> StringManager::fuzzy_wide_default_dict_;
std::vector<StringManager::FuzzyWideEntry> StringManager::fuzzy_wide_append_dict_;

// AC自动机缓存
StringManager::ACAutoCache StringManager::fuzzy_ac_auto_cache_default;
StringManager::ACAutoCache StringManager::fuzzy_ac_auto_cache_append;

// 锁
std::shared_mutex StringManager::rules_mutex_;
std::shared_mutex StringManager::fuzzy_mutex_;

// =================== AC自动机节点实现 ===================
struct StringManager::ACNodeUtf8 {
	std::map<char, std::unique_ptr<ACNodeUtf8>> children;
	ACNodeUtf8* fail = nullptr;
	struct MatchInfo {
		size_t rule_index;   // 规则索引
		size_t length;       // 匹配长度
	};
	std::vector<MatchInfo> outputs; // 输出规则集合
};

struct StringManager::ACNodeWide {
	std::map<wchar_t, std::unique_ptr<ACNodeWide>> children;
	ACNodeWide* fail = nullptr;
	struct MatchInfo {
		size_t rule_index;   // 规则索引
		size_t length;       // 匹配长度
	};
	std::vector<MatchInfo> outputs; // 输出规则集合
};

// =================== 基本初始化和清理 ===================
void StringManager::Initialize() {
	try {
		// 1. 初始化精确匹配规则表
		{
			std::unique_lock<std::shared_mutex> lock(rules_mutex_);
			default_utf8_rules_.reserve(std::size(DefaultRules));
			append_utf8_rules_.reserve(std::size(AppendRules));
			default_wide_rules_.reserve(std::size(WideDefaultRules));
			append_wide_rules_.reserve(std::size(WideAppendRules));

			for (const auto& rule : DefaultRules) {
				if (rule.first && rule.second) {
					default_utf8_rules_.emplace(rule.first, rule.second);
				}
			}
			for (const auto& rule : AppendRules) {
				if (rule.first && rule.second) {
					append_utf8_rules_.emplace(rule.first, rule.second);
				}
			}
			for (const auto& rule : WideDefaultRules) {
				if (rule.first && rule.second) {
					default_wide_rules_.emplace(rule.first, rule.second);
				}
			}
			for (const auto& rule : WideAppendRules) {
				if (rule.first && rule.second) {
					append_wide_rules_.emplace(rule.first, rule.second);
				}
			}
		}

		// 2. 初始化模糊匹配字典
		{
			std::unique_lock<std::shared_mutex> fuzzy_lock(fuzzy_mutex_);

			// 默认UTF-8规则
			for (const auto& rule : FuzzyUtf8DefaultRules) {
				if (std::get<0>(rule) && std::get<1>(rule)) {
					fuzzy_utf8_default_dict_.push_back({
						std::get<0>(rule),
						std::get<1>(rule),
						false,
						std::get<2>(rule)
						});
				}
			}

			// 追加UTF-8规则
			for (const auto& rule : FuzzyUtf8AppendRules) {
				if (std::get<0>(rule) && std::get<1>(rule)) {
					fuzzy_utf8_append_dict_.push_back({
						std::get<0>(rule),
						std::get<1>(rule),
						true,
						std::get<2>(rule)
						});
				}
			}

			// 默认宽字符规则
			for (const auto& rule : FuzzyWideDefaultRules) {
				if (std::get<0>(rule) && std::get<1>(rule)) {
					fuzzy_wide_default_dict_.push_back({
						std::get<0>(rule),
						std::get<1>(rule),
						false,
						std::get<2>(rule)
						});
				}
			}

			// 追加宽字符规则
			for (const auto& rule : FuzzyWideAppendRules) {
				if (std::get<0>(rule) && std::get<1>(rule)) {
					fuzzy_wide_append_dict_.push_back({
						std::get<0>(rule),
						std::get<1>(rule),
						true,
						std::get<2>(rule)
						});
				}
			}
		}

		// 3. 排序字典并清理AC自动机缓存
		SortFuzzyDicts();
	}
	catch (const std::exception& ex) {
		throw std::runtime_error(std::string("StringManager::Initialize failed: ") + ex.what());
	}
}

void StringManager::Cleanup() {
	// 清理精确匹配规则
	{
		std::unique_lock<std::shared_mutex> lock(rules_mutex_);
		default_utf8_rules_.clear();
		append_utf8_rules_.clear();
		default_wide_rules_.clear();
		append_wide_rules_.clear();
	}

	// 清理模糊匹配字典
	{
		std::unique_lock<std::shared_mutex> fuzzy_lock(fuzzy_mutex_);
		fuzzy_utf8_default_dict_.clear();
		fuzzy_utf8_append_dict_.clear();
		fuzzy_wide_default_dict_.clear();
		fuzzy_wide_append_dict_.clear();
	}

	// 清理AC自动机缓存
	ClearFuzzyAutomatonCache();
}

// =================== 规则增添 ===================
void StringManager::AddDefaultRule(const wchar_t* orig, const wchar_t* repl) {
	if (!orig || !repl) return;
	std::unique_lock<std::shared_mutex> lock(rules_mutex_);
	default_wide_rules_.emplace(orig, repl);
}

void StringManager::AddAppendRule(const wchar_t* orig, const wchar_t* repl) {
	if (!orig || !repl) return;
	std::unique_lock<std::shared_mutex> lock(rules_mutex_);
	append_wide_rules_.emplace(orig, repl);
}

void StringManager::AddDefaultRule(const char* orig, const char* repl) {
	if (!orig || !repl) return;
	std::unique_lock<std::shared_mutex> lock(rules_mutex_);
	default_utf8_rules_.emplace(orig, repl);
}

void StringManager::AddAppendRule(const char* orig, const char* repl) {
	if (!orig || !repl) return;
	std::unique_lock<std::shared_mutex> lock(rules_mutex_);
	append_utf8_rules_.emplace(orig, repl);
}

void StringManager::AddFuzzyDefaultRule(const wchar_t* orig, const wchar_t* repl, int priority) {
	if (!orig || !repl) return;
	std::unique_lock<std::shared_mutex> lock(fuzzy_mutex_);
	fuzzy_wide_default_dict_.push_back({ orig, repl, false, priority });
	SortFuzzyDicts();
}

void StringManager::AddFuzzyAppendRule(const wchar_t* orig, const wchar_t* repl, int priority) {
	if (!orig || !repl) return;
	std::unique_lock<std::shared_mutex> lock(fuzzy_mutex_);
	fuzzy_wide_append_dict_.push_back({ orig, repl, true, priority });
	SortFuzzyDicts();
}

void StringManager::AddFuzzyDefaultRule(const char* orig, const char* repl, int priority) {
	if (!orig || !repl) return;
	std::unique_lock<std::shared_mutex> lock(fuzzy_mutex_);
	fuzzy_utf8_default_dict_.push_back({ orig, repl, false, priority });
	SortFuzzyDicts();
}

void StringManager::AddFuzzyAppendRule(const char* orig, const char* repl, int priority) {
	if (!orig || !repl) return;
	std::unique_lock<std::shared_mutex> lock(fuzzy_mutex_);
	fuzzy_utf8_append_dict_.push_back({ orig, repl, true, priority });
	SortFuzzyDicts();
}

// =================== 精确查找 ===================
const char* StringManager::FindDefault(const char* orig) {
	if (!orig) return nullptr;
	std::shared_lock<std::shared_mutex> lock(rules_mutex_);
	auto it = default_utf8_rules_.find(orig);
	return it != default_utf8_rules_.end() ? it->second.c_str() : nullptr;
}

const char* StringManager::FindAppend(const char* orig) {
	if (!orig) return nullptr;
	std::shared_lock<std::shared_mutex> lock(rules_mutex_);
	auto it = append_utf8_rules_.find(orig);
	return it != append_utf8_rules_.end() ? it->second.c_str() : nullptr;
}

std::string StringManager::FindDefault(const std::string& orig) {
	std::shared_lock<std::shared_mutex> lock(rules_mutex_);
	auto it = default_utf8_rules_.find(orig);
	return it != default_utf8_rules_.end() ? it->second : "";
}

std::string StringManager::FindAppend(const std::string& orig) {
	std::shared_lock<std::shared_mutex> lock(rules_mutex_);
	auto it = append_utf8_rules_.find(orig);
	return it != append_utf8_rules_.end() ? it->second : "";
}

const wchar_t* StringManager::FindDefault(const wchar_t* orig) {
	if (!orig) return nullptr;
	std::shared_lock<std::shared_mutex> lock(rules_mutex_);
	auto it = default_wide_rules_.find(orig);
	return it != default_wide_rules_.end() ? it->second.c_str() : nullptr;
}

const wchar_t* StringManager::FindAppend(const wchar_t* orig) {
	if (!orig) return nullptr;
	std::shared_lock<std::shared_mutex> lock(rules_mutex_);
	auto it = append_wide_rules_.find(orig);
	return it != append_wide_rules_.end() ? it->second.c_str() : nullptr;
}

std::wstring StringManager::FindDefault(const std::wstring& orig) {
	std::shared_lock<std::shared_mutex> lock(rules_mutex_);
	auto it = default_wide_rules_.find(orig);
	return it != default_wide_rules_.end() ? it->second : L"";
}

std::wstring StringManager::FindAppend(const std::wstring& orig) {
	std::shared_lock<std::shared_mutex> lock(rules_mutex_);
	auto it = append_wide_rules_.find(orig);
	return it != append_wide_rules_.end() ? it->second : L"";
}

// =================== 模糊查找 ===================
std::string StringManager::FuzzyMatchDefault(const char* input) {
	if (!input) return "";
	return FuzzyMatchDefault(std::string(input));
}

std::string StringManager::FuzzyMatchAppend(const char* input) {
	if (!input) return "";
	return FuzzyMatchAppend(std::string(input));
}

std::string StringManager::FuzzyMatchDefault(const std::string& input) {
	std::shared_lock<std::shared_mutex> lock(fuzzy_mutex_);
	return PerformUtf8SubstringReplacement(input, fuzzy_utf8_default_dict_);
}

std::string StringManager::FuzzyMatchAppend(const std::string& input) {
	std::shared_lock<std::shared_mutex> lock(fuzzy_mutex_);
	return PerformUtf8SubstringReplacement(input, fuzzy_utf8_append_dict_);
}

std::wstring StringManager::FuzzyMatchDefault(const wchar_t* input) {
	if (!input) return L"";
	return FuzzyMatchDefault(std::wstring(input));
}

std::wstring StringManager::FuzzyMatchAppend(const wchar_t* input) {
	if (!input) return L"";
	return FuzzyMatchAppend(std::wstring(input));
}

std::wstring StringManager::FuzzyMatchDefault(const std::wstring& input) {
	std::shared_lock<std::shared_mutex> lock(fuzzy_mutex_);
	return PerformWideSubstringReplacement(input, fuzzy_wide_default_dict_);
}

std::wstring StringManager::FuzzyMatchAppend(const std::wstring& input) {
	std::shared_lock<std::shared_mutex> lock(fuzzy_mutex_);
	return PerformWideSubstringReplacement(input, fuzzy_wide_append_dict_);
}

// =================== 字典哈希与排序 ===================
size_t StringManager::ComputeFuzzyUtf8DictHash(const std::vector<FuzzyUtf8Entry>& dict) {
	size_t hash_value = dict.size();
	for (const auto& entry : dict) {
		hash_value ^= std::hash<std::string>()(entry.orig);
		hash_value ^= std::hash<std::string>()(entry.repl) << 1;
		hash_value ^= std::hash<int>()(entry.priority) << 2;
		hash_value ^= std::hash<bool>()(entry.is_append) << 3;
	}
	return hash_value;
}

size_t StringManager::ComputeFuzzyWideDictHash(const std::vector<FuzzyWideEntry>& dict) {
	size_t hash_value = dict.size();
	for (const auto& entry : dict) {
		hash_value ^= std::hash<std::wstring>()(entry.orig);
		hash_value ^= std::hash<std::wstring>()(entry.repl) << 1;
		hash_value ^= std::hash<int>()(entry.priority) << 2;
		hash_value ^= std::hash<bool>()(entry.is_append) << 3;
	}
	return hash_value;
}

void StringManager::SortFuzzyDicts() {
	// 排序函数：按优先级、字符串长度排序
	auto sort_fuzzy_dict = [](auto& dict) {
		std::sort(dict.begin(), dict.end(), [](const auto& a, const auto& b) {
			if (a.priority != b.priority)
				return a.priority > b.priority;
			return a.orig.size() > b.orig.size();
			});
		};

	// 分别排序四种字典
	sort_fuzzy_dict(fuzzy_utf8_default_dict_);
	sort_fuzzy_dict(fuzzy_utf8_append_dict_);
	sort_fuzzy_dict(fuzzy_wide_default_dict_);
	sort_fuzzy_dict(fuzzy_wide_append_dict_);

	ClearFuzzyAutomatonCache();
}

// =================== 清除AC自动机缓存 ===================
void StringManager::ClearFuzzyAutomatonCache() {
	// 清理默认规则缓存
	{
		std::lock_guard<std::mutex> lock(fuzzy_ac_auto_cache_default.mutex);
		fuzzy_ac_auto_cache_default.utf8_root.reset();
		fuzzy_ac_auto_cache_default.utf8_dict_hash = 0;
		fuzzy_ac_auto_cache_default.wide_root.reset();
		fuzzy_ac_auto_cache_default.wide_dict_hash = 0;
	}

	// 清理追加规则缓存
	{
		std::lock_guard<std::mutex> lock(fuzzy_ac_auto_cache_append.mutex);
		fuzzy_ac_auto_cache_append.utf8_root.reset();
		fuzzy_ac_auto_cache_append.utf8_dict_hash = 0;
		fuzzy_ac_auto_cache_append.wide_root.reset();
		fuzzy_ac_auto_cache_append.wide_dict_hash = 0;
	}
}

// =================== AC自动机构建和缓存 ===================
std::shared_ptr<StringManager::ACNodeUtf8> StringManager::BuildACAutomatonUtf8(
	const std::vector<FuzzyUtf8Entry>& dict)
{
	auto root = std::make_shared<ACNodeUtf8>();

	// 1. 构建Trie树
	for (size_t i = 0; i < dict.size(); ++i) {
		const auto& entry = dict[i];
		ACNodeUtf8* current = root.get();

		for (char c : entry.orig) {
			auto& next = current->children[c];
			if (!next) {
				next = std::make_unique<ACNodeUtf8>();
			}
			current = next.get();
		}

		// 在结束节点添加匹配信息
		current->outputs.push_back({
			i,          // 规则索引
			entry.orig.size() // 匹配长度
			});
	}

	// 2. 构建失败指针（BFS）
	std::queue<ACNodeUtf8*> q;

	// 根节点的失败指针指向自身
	root->fail = root.get();

	// 第一层节点的失败指针指向根节点
	for (auto& [c, child] : root->children) {
		child->fail = root.get();
		q.push(child.get());
	}

	while (!q.empty()) {
		ACNodeUtf8* current = q.front();
		q.pop();

		for (auto& [c, child] : current->children) {
			ACNodeUtf8* fail = current->fail;

			// 沿着失败指针向上查找，直到找到匹配的子节点或到达根节点
			while (fail != root.get() && fail->children.find(c) == fail->children.end()) {
				fail = fail->fail;
			}

			// 如果找到匹配的子节点，设置失败指针
			if (fail->children.find(c) != fail->children.end()) {
				child->fail = fail->children[c].get();
			}
			else {
				child->fail = root.get();
			}

			// 合并输出
			for (const auto& output : child->fail->outputs) {
				child->outputs.push_back(output);
			}

			q.push(child.get());
		}
	}

	return root;
}

std::shared_ptr<StringManager::ACNodeWide> StringManager::BuildACAutomatonWide(
	const std::vector<FuzzyWideEntry>& dict)
{
	auto root = std::make_shared<ACNodeWide>();

	// 1. 构建Trie树
	for (size_t i = 0; i < dict.size(); ++i) {
		const auto& entry = dict[i];
		ACNodeWide* current = root.get();

		for (wchar_t c : entry.orig) {
			auto& next = current->children[c];
			if (!next) {
				next = std::make_unique<ACNodeWide>();
			}
			current = next.get();
		}

		// 在结束节点添加匹配信息
		current->outputs.push_back({
			i,          // 规则索引
			entry.orig.size() // 匹配长度
			});
	}

	// 2. 构建失败指针（BFS）
	std::queue<ACNodeWide*> q;

	// 根节点的失败指针指向自身
	root->fail = root.get();

	// 第一层节点的失败指针指向根节点
	for (auto& [c, child] : root->children) {
		child->fail = root.get();
		q.push(child.get());
	}

	while (!q.empty()) {
		ACNodeWide* current = q.front();
		q.pop();

		for (auto& [c, child] : current->children) {
			ACNodeWide* fail = current->fail;

			// 沿着失败指针向上查找，直到找到匹配的子节点或到达根节点
			while (fail != root.get() && fail->children.find(c) == fail->children.end()) {
				fail = fail->fail;
			}

			// 如果找到匹配的子节点，设置失败指针
			if (fail->children.find(c) != fail->children.end()) {
				child->fail = fail->children[c].get();
			}
			else {
				child->fail = root.get();
			}

			// 合并输出
			for (const auto& output : child->fail->outputs) {
				child->outputs.push_back(output);
			}

			q.push(child.get());
		}
	}

	return root;
}
// =================== AC自动机匹配函数 ===================
std::vector<std::pair<size_t, size_t>> StringManager::MatchUsingACAutomaton(
	const std::shared_ptr<ACNodeUtf8>& root,
	const std::string& input)
{
	std::vector<std::pair<size_t, size_t>> matches;
	ACNodeUtf8* current = root.get();

	for (size_t i = 0; i < input.size(); ++i) {
		char c = input[i];

		// 沿着失败指针向上查找，直到找到匹配的子节点或到达根节点
		while (current != root.get() && current->children.find(c) == current->children.end()) {
			current = current->fail;
		}

		// 如果找到匹配的子节点，移动到该节点
		if (current->children.find(c) != current->children.end()) {
			current = current->children[c].get();

			// 收集所有匹配
			for (const auto& output : current->outputs) {
				size_t start = i - output.length + 1;
				matches.emplace_back(start, output.rule_index);
			}
		}
	}

	return matches;
}

std::vector<std::pair<size_t, size_t>> StringManager::MatchUsingACAutomaton(
	const std::shared_ptr<ACNodeWide>& root,
	const std::wstring& input)
{
	std::vector<std::pair<size_t, size_t>> matches;
	ACNodeWide* current = root.get();

	for (size_t i = 0; i < input.size(); ++i) {
		wchar_t c = input[i];

		// 沿着失败指针向上查找，直到找到匹配的子节点或到达根节点
		while (current != root.get() && current->children.find(c) == current->children.end()) {
			current = current->fail;
		}

		// 如果找到匹配的子节点，移动到该节点
		if (current->children.find(c) != current->children.end()) {
			current = current->children[c].get();

			// 收集所有匹配
			for (const auto& output : current->outputs) {
				size_t start = i - output.length + 1;
				matches.emplace_back(start, output.rule_index);
			}
		}
	}

	return matches;
}


// =================== AC自动机支持的替换 ===================
std::string StringManager::PerformUtf8SubstringReplacement(
	const std::string& input,
	const std::vector<FuzzyUtf8Entry>& dict)
{
	if (dict.empty()) return input;

	// 获取合适的缓存
	auto& cache = dict[0].is_append
		? fuzzy_ac_auto_cache_append
		: fuzzy_ac_auto_cache_default;

	// 计算字典哈希
	size_t dict_hash = ComputeFuzzyUtf8DictHash(dict);

	// 检查缓存
	std::shared_ptr<ACNodeUtf8> root;
	{
		std::lock_guard<std::mutex> lock(cache.mutex);
		if (!cache.utf8_root || cache.utf8_dict_hash != dict_hash) {
			cache.utf8_root = BuildACAutomatonUtf8(dict);
			cache.utf8_dict_hash = dict_hash;
		}
		root = cache.utf8_root;
	}

	// 使用AC自动机进行匹配
	auto matches = MatchUsingACAutomaton(root, input);

	// 如果没有匹配，直接返回输入
	if (matches.empty()) return input;

	// 对匹配进行排序（按起始位置）
	std::sort(matches.begin(), matches.end(),
		[](const auto& a, const auto& b) {
			return a.first < b.first;
		});

	// 应用替换（避免重叠）
	std::string result;
	size_t last_pos = 0;
	size_t current_end = 0;

	for (const auto& match : matches) {
		size_t start = match.first;
		size_t rule_index = match.second;

		// 跳过重叠的匹配
		if (start < current_end) continue;

		// 添加未匹配部分
		result.append(input.substr(last_pos, start - last_pos));

		// 添加替换字符串
		result.append(dict[rule_index].repl);

		// 更新位置
		last_pos = start + dict[rule_index].orig.size();
		current_end = last_pos;
	}

	// 添加剩余部分
	if (last_pos < input.size()) {
		result.append(input.substr(last_pos));
	}

	return result;
}

std::wstring StringManager::PerformWideSubstringReplacement(
	const std::wstring& input,
	const std::vector<FuzzyWideEntry>& dict)
{
	if (dict.empty()) return input;

	// 获取合适的缓存
	auto& cache = dict[0].is_append
		? fuzzy_ac_auto_cache_append
		: fuzzy_ac_auto_cache_default;

	// 计算字典哈希
	size_t dict_hash = ComputeFuzzyWideDictHash(dict);

	// 检查缓存
	std::shared_ptr<ACNodeWide> root;
	{
		std::lock_guard<std::mutex> lock(cache.mutex);
		if (!cache.wide_root || cache.wide_dict_hash != dict_hash) {
			cache.wide_root = BuildACAutomatonWide(dict);
			cache.wide_dict_hash = dict_hash;
		}
		root = cache.wide_root;
	}

	// 使用AC自动机进行匹配
	auto matches = MatchUsingACAutomaton(root, input);

	// 如果没有匹配，直接返回输入
	if (matches.empty()) return input;

	// 对匹配进行排序（按起始位置）
	std::sort(matches.begin(), matches.end(),
		[](const auto& a, const auto& b) {
			return a.first < b.first;
		});

	// 应用替换（避免重叠）
	std::wstring result;
	size_t last_pos = 0;
	size_t current_end = 0;

	for (const auto& match : matches) {
		size_t start = match.first;
		size_t rule_index = match.second;

		// 跳过重叠的匹配
		if (start < current_end) continue;

		// 添加未匹配部分
		result.append(input.substr(last_pos, start - last_pos));

		// 添加替换字符串
		result.append(dict[rule_index].repl);

		// 更新位置
		last_pos = start + dict[rule_index].orig.size();
		current_end = last_pos;
	}

	// 添加剩余部分
	if (last_pos < input.size()) {
		result.append(input.substr(last_pos));
	}

	return result;
}